package tx

import java.util.concurrent.locks.ReentrantReadWriteLock
import scala.collection.JavaConversions

/** Transactional Extender: STM that wraps object in transactions. */
class Tx {
    class RetryTransaction(s: String) extends RuntimeException(s)
    val transactions = new ThreadLocal[Transaction] { override def initialValue() = new Transaction() }

    def TLock() = new TLock()
    class TLock extends ReentrantReadWriteLock {
        protected[Tx] var version: Long = 1;
        
        def rLock(oldVersion: Long) = {
            this.readLock.lock
            if (oldVersion == 0 || oldVersion == version) version
            else throw new RetryTransaction("Object changed since first access.")
        }
        def rUnlock() = this.readLock.unlock

        def wLock(oldVersion: Long) = {
            this.writeLock.lock
            if (oldVersion != 0 && oldVersion != version) {
                this.writeLock.unlock
                throw new RetryTransaction("Object changed since first access.")
            } else version // no lock upgrade even required. damn.
        }
        def wUnlock() = { val temp = version; this.writeLock.unlock; temp }
    }

    class Transaction {
        val reads = JavaConversions.asMap(new java.util.HashMap[TLock, Long]())
        val writes = JavaConversions.asSet(new java.util.HashSet[TLock]())

        def withReadLock[E](lock: TLock)(action: => E): E = {
            if (writes.contains(lock)) action else { //write-locked, nothing more to do.
                try reads.get(lock) match {
                    case Some(version) => { lock.rLock(version); action } //revalidate version.
                    case None => { reads(lock) = lock.rLock(0); action }
                } finally lock.rUnlock() //after performing action
            }
        }

        def withWriteLock[E](lock: TLock)(action: => E): E = {
            if (writes.contains(lock)) return action else { //already write-locked!
                reads.get(lock) match {
                    case None => { lock.wLock(0); action } //lock and act
                    case Some(version) => { 
                        reads -= lock; // drop the read "lock"
                        lock.wLock(version); writes += lock // validate version, acquire write lock.
                        action // now we can do the deed.
                    }
                }
            }
        }

        def prepare() {

        }
    }

    class Ref[E](initial: E) {
        val lock = new TLock
        var current = initial;

        def apply(implicit t: Transaction) = t.withReadLock(lock) { current }
        def update (value: E)(implicit t:Transaction) = t.withWriteLock(lock) { current = value }
    }
}

object Tx extends Tx